import { ImageGenerationParams } from "../types/index.js";
import { replicate } from "../services/replicate.js";
import { handleError } from "../utils/error.js";
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { FileOutput } from "replicate";
import { outputToBase64 } from "../utils/image.js";
import { CONFIG } from "../config/index.js";

export const registerGenerateImageTool = async (
  input: ImageGenerationParams
): Promise<CallToolResult> => {
  const { support_image_mcp_response_type, ...predictionInput } = input;
  try {
    const [output] = (await replicate.run(CONFIG.imageModelId, {
      input: predictionInput,
    })) as [FileOutput];
    const imageUrl = output.url() as unknown as string;

    if (support_image_mcp_response_type) {
      const imageBase64 = await outputToBase64(output);
      return {
        content: [
          {
            type: "text",
            text: `This is a generated image link: ${imageUrl}`,
          },
          {
            type: "image",
            data: imageBase64,
            mimeType: "image/png",
          },
          {
            type: "text",
            text: `The image above is generated by the Flux model and prompt: ${input.prompt}`,
          },
        ],
      };
    }

    return {
      content: [
        {
          type: "text",
          text: `This is a generated image link: ${imageUrl}`,
        },
        {
          type: "text",
          text: `The image above is generated by the Flux model and prompt: ${input.prompt}`,
        },
      ],
    };
  } catch (error) {
    handleError(error);
  }
};
